# src/plot_data/other/communication.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*

import os
import json
import argparse
import sys 
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

sys.path.append("../../")
from package.utility import ensure_exists
from package.utility import find_targets
sys.path.append("..")
from baseplot import casename2number
from baseplot import create_subplot
from baseplot import mypaper_subplot
from baseplot import close_sublot_levelplot
from baseplot import mypaper_subplot_level
from baseplot import mypaper_subplot_level_axset

# set parser
def make_parser(): 
    parser = argparse.ArgumentParser(
        prog="jxpamg-predo-merge-pid-timerdata",
        description="jxpamg timer data preprocess script, merge pid's timer data to a single json.")
    parser.add_argument(
            "-sdir","--savedir",
            dest="savedir",
            default="./",
            help="Destination save pic/csv directories."
        )
    # parser.add_argument(
    #         "-po", "--picopt",
    #         dest="plot_type",
    #         default="time_linechart",
    #         help="Set type of pic out. Option: rank, time_heatmap, time_linechart, percent_par, percent_total, percent_line, efficiency, sizeperprocess."
    # )
    # parser.add_argument(
    #         "-co", "--csvopt",
    #         dest="csv_type",
    #         default="time",
    #         help="Set type of data in csv. Option: time, rank_correlation."
    # )
    # parser.add_argument(
    #         "-o", "--output",
    #         dest="out_type",
    #         default="csv",
    #         help="Set output type of plotting picture. Option: csv, pic."
    # )
    parser.add_argument(
            "LOGDIR",
            metavar="LOGDIR",
            help="Destination log directories."
        )
    # parser.add_argument(
    #         "CASENAME",
    #         metavar="CASENAME",
    #         default="jxpamg_solver01",
    #         nargs='+',
    #         help="Set casename of plotting picture. Option: jxpamg_lixue_matrix, jxpamg_shangfei_iter00001_job0, jxpamg_solver01, jxpamg_solver02, jxpamg_solver03, jxpamg_solver04."
    # )
    return parser

def get_runcases(topdir,jsonname="*-compute_length.json"):
    """Get runcase under topdir."""
    topdir = os.path.abspath(topdir)
    ensure_exists(topdir)

    print(f">> scan runcase under {topdir}") 
    runcases = find_targets(topdir, jsonname, suffix="-xtype f") 
    runcase_num = len(runcases)
    print(f">> {runcase_num} runcases found!")

    return runcases

def getlength(dict):
    lengthlist = []
    for i in dict:
        np = i.split("_")[-1]
        for j in dict[i]:
            tmpdict = {}
            level = j.split("_")[-1]
            tmpdict["np"] = int(np)
            tmpdict["level"] = int(level)
            for k in dict[i][j]:
                tmpdict[k] = float(dict[i][j][k])
            lengthlist.append(tmpdict)
    df = pd.DataFrame(lengthlist)
    return df

# 获得不同 case 得到 component 对应的 time
def getjsondata(casepath):
    timedic = {}
    for i in casepath:
        basename = os.path.basename(i)
        basename = basename.split("-")[0]
        with open(i, "r") as fo:
            text = fo.read()
            jsondata = json.loads(text)
            print(f"Load {basename} data.")
        timedic[basename] = getlength(jsondata)
    return timedic

def myplot(savedir,dict):
    for k in ["ave","max"]:
        fig, axes = create_subplot(nrows=2, ncols=3, xlable="Number  of  Processors", ylable="Communicating Range")
        legend_name = []
        for i in dict:
            caseid = casename2number(i)
            casenumber = int(caseid.split("C")[-1]) - 1
            df = dict[i]
            levellist = list(df["level"])
            levellist = sorted(set(levellist), key=levellist.index)
            plotdf = df[["np","level",k]]
            # axes[casenumber//3, casenumber%3].set_yscale("log", nonpositive='clip')

            title = f"{caseid}_{k}-Range"
            for j in levellist:
                if j not in legend_name:
                    legend_name.append(f"Level {j}")
                levelplotdf = plotdf[plotdf["level"]==j].sort_values(by=["np"])
                levelplotdf = levelplotdf.drop(["level"], axis=1)
                levelplotdf.set_index(["np"], inplace=True)
                mypaper_subplot_level(axes[casenumber//3, casenumber%3], levelplotdf)
            mypaper_subplot_level_axset(axes[casenumber//3, casenumber%3], title, levelplotdf)
        close_sublot_levelplot(2,3,fig,axes,savedir,title,"line",legend_name)

def mycsv(savedir,dict):
    for i in dict:
        caseid = casename2number(i)
        df = dict[i]
        levellist = list(df["level"])
        levellist = sorted(list(set(levellist)))
        for k in ["ave","max"]:
            plotdf = df[["np","level",k]]
            csvdf_list = []
            for j in levellist:
                levelplotdf = plotdf[plotdf["level"]==j].sort_values(by=["np"])
                levelplotdf = levelplotdf.drop(["level"], axis=1)
                # csvdf = levelplotdf.set_index('np') 
                levelplotdf.rename(columns={f'{k}':f"{k}_l{j}"}, inplace = True)
                csvdf_list.append(levelplotdf)
            for j in range(len(csvdf_list)):
                if(j == 0):
                    csvdf = csvdf_list[j]
                else:
                    csvdf = pd.merge(csvdf,csvdf_list[j],how='outer')
            csvdf.to_csv(os.path.join(savedir, f"{caseid}-{k}-length.csv"), index=False)
            # writer = pd.ExcelWriter(os.path.join(savedir, f"c{caseid}-{k}-length.xlsx"))
            # csvdf.to_excel(writer,sheet_name=k)
            # writer.save()

def main():
    parser = make_parser()
    args = parser.parse_args()
    runcases = get_runcases(args.LOGDIR)
    # print(runcases)
    datadict = getjsondata(runcases)
    # print(list(datadict["jxpamg_solver04"]["np"]))
    myplot(args.savedir,datadict)
    # mycsv(args.savedir,datadict)

if __name__ == "__main__":
    main()